package com.example.dynamicdatasource.config.datasource;

import com.baomidou.dynamic.datasource.DynamicRoutingDataSource;
import com.baomidou.dynamic.datasource.creator.DataSourceCreator;
import com.baomidou.dynamic.datasource.ds.GroupDataSource;
import com.baomidou.dynamic.datasource.exception.CannotFindDataSourceException;
import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.DataSourceProperty;
import com.example.dynamicdatasource.config.ProvinceDbProperties;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;


@Slf4j
public class LazyDynamicDataSource extends DynamicRoutingDataSource {

    private final ProvinceDbProperties provinceDbProperties;
    private final DataSourceCreator dataSourceCreator;
    private final String urlPattern;
    private final String username;
    private final String password;
    private final String driverClassName;

    // 数据源最后使用时间记录
    private final Map<String, Long> dataSourceLastUsedTime = new ConcurrentHashMap<>();
    private final ReentrantLock lock = new ReentrantLock();



    public LazyDynamicDataSource(ProvinceDbProperties provinceDbProperties,
                                DataSourceCreator dataSourceCreator,
                                String urlPattern,
                                String username,
                                String password,
                                String driverClassName,
                                Map<String, DataSource> dataSources) {
        this.provinceDbProperties = provinceDbProperties;
        this.dataSourceCreator = dataSourceCreator;
        this.urlPattern = urlPattern;
        this.username = username;
        this.password = password;
        this.driverClassName = driverClassName;

        dataSources.forEach((key, value) -> {
            addDataSource(key, value);
        });
    }

    @Override
    public DataSource determineDataSource() {
        String dsKey = DynamicDataSourceContextHolder.peek();

        // 如果是主数据源或已加载的数据源，直接返回
        if (!provinceDbProperties.getConfigs().containsKey(dsKey) || getDataSource(dsKey) != null) {
            updateLastUsedTime(dsKey);
            return super.determineDataSource();
        }

        // 检查是否需要懒加载
        ProvinceDbProperties.ProvinceDbConfig config = provinceDbProperties.getConfigs().get(dsKey);
        if (config == null) {
            log.warn("未找到数据源配置: {}", dsKey);
            updateLastUsedTime(dsKey);
            return super.determineDataSource();
        }

        // 懒加载数据源（双重检查锁确保线程安全）
        try {
            lock.lock();
            if (getDataSource(dsKey) == null) {
                log.info("懒加载数据源: {}", dsKey);
                String url = urlPattern.replace("{dbName}", config.getName());

                // 构建数据源属性对象（修正部分）
                DataSourceProperty property = new DataSourceProperty();
                property.setPoolName(dsKey);
                property.setUrl(url);
                property.setUsername(username);
                property.setPassword(password);
                property.setDriverClassName(driverClassName);

                DataSource dataSource = dataSourceCreator.createDataSource(property);
                addDataSource(dsKey, dataSource);
            }
        } finally {
            lock.unlock();
        }

        updateLastUsedTime(dsKey);
        return super.determineDataSource();
    }

    @Override
    public synchronized void addDataSource(String ds, DataSource dataSource) {
        updateLastUsedTime(ds);
        super.addDataSource(ds, dataSource);
    }

    @Override
    public DataSource getDataSource(String ds) {
        return super.getDataSources().getOrDefault(ds, null);
    }

    /**
     * 更新数据源最后使用时间
     */
    public void updateLastUsedTime(String dsKey) {
        // 不记录主数据源的使用时间
        if (!provinceDbProperties.getConfigs().containsKey(dsKey)) {
            return;
        }
        dataSourceLastUsedTime.put(dsKey, System.currentTimeMillis());
    }

    /**
     * 获取数据源最后使用时间
     */
    public Long getLastUsedTime(String dsKey) {
        return dataSourceLastUsedTime.get(dsKey);
    }

    /**
     * 关闭并移除超时的数据源
     */
    public void closeTimeoutDataSources(long timeoutMillis) {
        long currentTime = System.currentTimeMillis();

        for (String dsKey : provinceDbProperties.getConfigs().keySet()) {
            DataSource dataSource = getDataSource(dsKey);
            if (dataSource == null) {
                continue;
            }

            Long lastUsedTime = getLastUsedTime(dsKey);
            if (lastUsedTime == null) {
                continue;
            }

            // 检查是否超时
            if (currentTime - lastUsedTime > timeoutMillis) {
                try {
                    lock.lock();
                    // 再次检查，避免并发问题
                    if (getDataSource(dsKey) != null) {
                        log.info("关闭超时数据源: {}", dsKey);
                        // 关闭数据源
                        if (dataSource instanceof AutoCloseable) {
                            ((AutoCloseable) dataSource).close();
                        }
                        // 从动态数据源中移除
                        removeDataSource(dsKey);
                        dataSourceLastUsedTime.remove(dsKey);
                    }
                } catch (Exception e) {
                    log.error("关闭数据源 {} 失败", dsKey, e);
                } finally {
                    lock.unlock();
                }
            }
        }
    }
}
